# coding: utf-8

import os
import time
import argparse
import model_config



def main():
    parser = argparse.ArgumentParser(description="Auto locust test")
    parser.add_argument("--url", type=str, default=model_config.url, help="模型地址url")
    parser.add_argument("--file", type=str, default=model_config.data_file, help="数据集，默认为yb_test_data.ndjson")
    parser.add_argument("--model_name", type=str, default=model_config.model_name, help="模型名称")
    parser.add_argument("--share_data", type=int, default=model_config.share_data,
                        help="不同并发是否共享数据，0：不共享，1：共享, 2: 共享且只跑一轮")
    parser.add_argument("--sample_type", type=int, default=model_config.sample_type, help="数据采样规则，0：顺序，1：随机")
    parser.add_argument("--api_key", type=str, default="", help="api key，可空")
    parser.add_argument("--sleep_time", type=int, default=5, help="api key，可空")
    parser.add_argument("--run_time", type=str, default='360m', help="api key，可空")

    cmd = 'locust -f client.py --headless -u %s -r 1 -t %s --csv=report/report_%s_%s --only-summary ' \
          '--url=%s --file=%s --model_name=%s --share_data=%s --sample_type=%s --api_key=%s '

    args = parser.parse_args()
    task_time = time.strftime("%Y-%m-%d_%H-%M-%S")
    for concurrency_num in [32, 64]:
    #for concurrency_num in [15, 40]:
        sys_cmd = cmd % (concurrency_num, args.run_time, task_time, "%02d" % concurrency_num, args.url, args.file, args.model_name,
                         args.share_data, args.sample_type, args.api_key)

        log_file = f'logs/{task_time}_b{"%02d" % concurrency_num}.log'
        sys_cmd = f'{sys_cmd} > {log_file} 2>&1'
        os.system("mkdir -p logs")
        print(sys_cmd)
        os.system(sys_cmd)
        if concurrency_num != 40:
            time.sleep(args.sleep_time)

    os.system('cat report/report_%s_*_stats.csv > report/report_stats_%s.csv' % (task_time, task_time))

if __name__ == '__main__':
    main()
